"""
- 输入图像
  - OpenCV打开图像的形状是HWC, 而网络输入的形状是NCHW, 需要进行通道转换。
  - 输入网络的数据进行了归一化操作: 数据比较大容易出现数值溢出；针对不同输入图像网络不需要去关心大小, 更多的关注特征本身。
  - 通常我们在计算的时候使用的numpy数组进行计算。比如说OpenCV打开图像 做一些数学运算的时候 使用都是Numpy数组，而网络输入的是Tensor，张量
    需要将Numpy数组转换成tensor
- 特征提取
  - 提取的特征我们将tensor转换为numpy数组
"""
import json
import cv2
import numpy as np
import torch
from face_detect import FaceDetect
from net.facenet import InceptionResnetV1

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')


class FaceFeat:
    def __init__(self):
        self.face_net = InceptionResnetV1(is_train=False).to(device)
        # 开启验证
        self.face_net.eval()
        # 加载参数
        self.face_net.load_state_dict(torch.load('./weights/facenet_best_server.pt'))
        self.db_names, self.db_feats = self.load_db_face_feats()

    def preprocess(self, face_img):
        """
        图像处理: 大小设置、通道调换、归一化
        """
        face_img = face_img.astype(np.float32)
        face_img = cv2.resize(face_img, dsize=(112, 112))
        half_value = np.max(face_img) / 2
        face_img = (face_img - half_value) / half_value
        face_img = np.transpose(face_img, (2, 0, 1))
        face_img = np.expand_dims(face_img, 0)
        return face_img

    def get_face_feat(self, face_img):
        """
        特征提取
        :param face_img 人脸图像
        :return face_feat 人脸特征
        """
        pre_face_img = self.preprocess(face_img)
        # 将numpy数组转换为tensor
        # 必须和face_net网络的参数在同一个设备
        img_tensor = torch.from_numpy(pre_face_img).to(device)
        feat = self.face_net(img_tensor)
        # 将数据转换到cpu上再进行转换为numpy数组
        feat = feat.detach().cpu().numpy()
        return feat[0]



    def load_db_face_feats(self, path='./files/feat'):
        """
        加载注册的人脸特征
        """
        db_names = []
        db_feats = []
        with open('./files/feat', 'r', encoding='utf-8') as f:
            lines = f.readlines()
            for line in lines:
                feat_id, name, feat_str = line.split("|")
                # 将json字符串转换为python对象
                feat = json.loads(feat_str)
                db_names.append(name)
                #相识度的比对
                db_feats.append(feat)
            db_feats = np.array(db_feats)
            return db_names, db_feats







    def cal_similarity(self, face_img):
        """
        :param face_feat 识别的人脸特征
        思路:
        加载注册的人脸特征 db_face_feats
        """
        #目标人脸特征
        target_feat = self.get_face_feat(face_img)
        #已存储人脸特征
        db_names, db_feats = self.db_names, self.db_feats
        # 范式
        # (128)
        # (48, 128)--->(48,)
        # 沿着那个轴就将那个轴进行折叠
        dist = np.linalg.norm(target_feat - db_feats, axis=1)
        stack_dist_names = np.column_stack((np.array(dist, dtype=object),
                                            db_names))
        # 按照距离排序
        idx = np.argsort(stack_dist_names[:, 0])
        # 排序之后的数据
        """
        [0.6003428421031193 'HX01']
        [0.6683060754928192 'HX01']
        [0.6781350150078158 'HX'] """
        stack_dist_names = stack_dist_names[idx][:3]
        pred_name = stack_dist_names[0, 1]
        cnt = np.count_nonzero(pred_name == stack_dist_names[:, 1])
        if cnt < 2:
            return "unknown"
        return pred_name




if __name__ == '__main__':
    face_feat = FaceFeat()

    face_img = cv2.imread('./img/hx.png')

    pred_name = face_feat.cal_similarity(face_img)

    pass
